package com.ruoyi.framework.config;

import org.apache.commons.lang3.StringUtils;
import org.slf4j.MDC;

import java.util.Collection;
import java.util.List;
import java.util.Map;
import java.util.concurrent.*;

/**
 * @desc MDC线程池
 */
@SuppressWarnings("all")
public class MDCThreadPoolExecutor extends ThreadPoolExecutor {
    public MDCThreadPoolExecutor(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue, ThreadFactory threadFactory, RejectedExecutionHandler handler) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, threadFactory, handler);
    }

    public MDCThreadPoolExecutor(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, Executors.defaultThreadFactory(), new AbortPolicy());
    }

    /**
     * 单线程的线程池 （代码中大量使用此类线程池，本次适配traceId，后续全部干掉这类线程池）
     *
     * @return
     */
    public static ExecutorService newSingleThreadExecutor() {
        return new FinalizableDelegatedExecutorService(new MDCThreadPoolExecutor(1, 1, 0L, TimeUnit.MILLISECONDS, new LinkedBlockingQueue<Runnable>()));
    }

    @Override
    public void execute(Runnable runnable) {
        if (runnable == null) {
            throw new NullPointerException();
        }
        super.execute(wrap(runnable));
    }

    @Override
    public Future<?> submit(Runnable task) {
        if (task == null) {
            throw new NullPointerException();
        }
        return super.submit(wrap(task));
    }

    @Override
    public <T> Future<T> submit(Callable<T> task) {
        if (task == null) {
            throw new NullPointerException();
        }
        return super.submit(wrap(task));
    }

    @Override
    public <T> Future<T> submit(Runnable task, T result) {
        if (task == null) {
            throw new NullPointerException();
        }
        return super.submit(wrap(task), result);
    }

    /**
     * 包装Callable（支持父子线程上下文传递与子线程上下文销毁）
     *
     * @param task
     * @param <T>
     * @return
     */
    private <T> Callable<T> wrap(Callable<T> task) {
        Map<String, String> copyOfContextMap = MDC.getCopyOfContextMap();
        //主线程名称
        String currentThreadName = Thread.currentThread().getName();
        return () -> {
            try {
                copyContextMap(copyOfContextMap);
                return task.call();
            } finally {
                clearMDC(currentThreadName);
            }
        };
    }

    /**
     * 包装Runnable（支持父子线程上下文传递与子线程上下文销毁）
     *
     * @param runnable
     * @return
     */
    private Runnable wrap(Runnable runnable) {
        Map<String, String> copyOfContextMap = MDC.getCopyOfContextMap();
        //主线程名称
        String currentThreadName = Thread.currentThread().getName();
        return () -> {
            try {
                copyContextMap(copyOfContextMap);
                runnable.run();
            } finally {
                // 线程池使用拒绝策略为CallerRunsPolicy时，避免提前清除主线程中的MDC上下文。
                // 主线程名称 与 子线程名称不同说明使用的是线程池中的线程，而不是主线程：上下文可以及时清除。
                clearMDC(currentThreadName);
            }
        };
    }

    /**
     * 复制主线程中的上下文到子线程
     *
     * @param copyOfContextMap
     */
    private void copyContextMap(Map<String, String> copyOfContextMap) {
        if (copyOfContextMap != null) {
            MDC.setContextMap(copyOfContextMap);
        }
    }

    /**
     * 清除子线程中的上下文
     *
     * @param currentThreadName
     */
    private void clearMDC(String currentThreadName) {
        // 线程池使用拒绝策略为CallerRunsPolicy时，避免提前清除主线程中的MDC上下文。
        // 主线程名称 与 子线程名称不同说明使用的是线程池中的线程，而不是主线程：上下文可以及时清除。
        if (!StringUtils.equals(currentThreadName, Thread.currentThread().getName())) {
            MDC.clear();
        }
    }

    // 以下为兼容项目中的单线程线程池，复制自Executors.FinalizableDelegatedExecutorService，后续逐步废弃。 //

    static class FinalizableDelegatedExecutorService extends DelegatedExecutorService {
        FinalizableDelegatedExecutorService(ExecutorService executor) {
            super(executor);
        }

        protected void finalize() {
            super.shutdown();
        }
    }

    static class DelegatedExecutorService extends AbstractExecutorService {
        private final ExecutorService e;

        DelegatedExecutorService(ExecutorService executor) {
            e = executor;
        }

        public void execute(Runnable command) {
            e.execute(command);
        }

        public void shutdown() {
            e.shutdown();
        }

        public List<Runnable> shutdownNow() {
            return e.shutdownNow();
        }

        public boolean isShutdown() {
            return e.isShutdown();
        }

        public boolean isTerminated() {
            return e.isTerminated();
        }

        public boolean awaitTermination(long timeout, TimeUnit unit) throws InterruptedException {
            return e.awaitTermination(timeout, unit);
        }

        public Future<?> submit(Runnable task) {
            return e.submit(task);
        }

        public <T> Future<T> submit(Callable<T> task) {
            return e.submit(task);
        }

        public <T> Future<T> submit(Runnable task, T result) {
            return e.submit(task, result);
        }

        public <T> List<Future<T>> invokeAll(Collection<? extends Callable<T>> tasks) throws InterruptedException {
            return e.invokeAll(tasks);
        }

        public <T> List<Future<T>> invokeAll(Collection<? extends Callable<T>> tasks, long timeout, TimeUnit unit) throws InterruptedException {
            return e.invokeAll(tasks, timeout, unit);
        }

        public <T> T invokeAny(Collection<? extends Callable<T>> tasks) throws InterruptedException, ExecutionException {
            return e.invokeAny(tasks);
        }

        public <T> T invokeAny(Collection<? extends Callable<T>> tasks, long timeout, TimeUnit unit) throws InterruptedException, ExecutionException, TimeoutException {
            return e.invokeAny(tasks, timeout, unit);
        }
    }

}